from base64 import decode
import numpy as np
from mindx.sdk import base
from mindx.sdk.base import Tensor, Model, Size, log, ImageProcessor, post, BTensor
import cv2

device_id = 0  # 芯片ID
model_path = "yolov3_tf_bs1_fp16.om"  # 模型的路径
config_path = "yolov3_tf_bs1_fp16.cfg"  # 模型配置文件的路径
label_path = "yolov3.names"  # 分类标签文件的路径
image_path = "test.jpg"  # 输入图片
b_usedvpp = True  # 使用dvpp图像处理器启用，使用opencv时False
yolo_resizelen = 416


def main():
    base.mx_init() # 全局资源初始化
    yolov3 = Model(model_path, device_id)  # 创造模型对象
    imageTensorList = []
    if b_usedvpp:
        # 创造图像处理器对象!!!!!使用该方法处理后数据已在device侧
        imageProcessor0 = ImageProcessor(device_id)
        decodedImg = imageProcessor0.decode(image_path, base.nv12)

        imageProcessor1 = ImageProcessor(device_id)
        size_cof = Size(yolo_resizelen, yolo_resizelen)
        resizeImg = imageProcessor1.resize(decodedImg, size_cof)

        imageTensorList = [resizeImg.to_tensor()]  # 推理前需要转换为tensor的List，数据已在device侧无需转移
    else:
        # 本示例中模型为使用了aipp的YOLOV3模型（输入为YUV格式），直接使用opencv解码无法正常运行，该处仅演示使用opencv进行图像处理和组装Tensor的方法
        image = np.array(cv2.imread(image_path))
        size_cof = (yolo_resizelen, yolo_resizelen)
        resizeImg = cv2.resize(image, size_cof, interpolation=cv2.INTER_LINEAR)

        imageTensor = Tensor(resizeImg) # 推理前需要转换为tensor的List，使用Tensor类来构建。
        imageTensor.to_device(device_id) # !!!!!重要，需要转移至device侧，该函数单独执行
        imageTensorList = [imageTensor] # 推理前需要转换为tensor的List

        """
        使用外部数据作为tensor时务必使用to_device进行转移，缺失该步骤会导致输出结果异常，RC3以上版本已修复
        ！！！如使用了transpose,slice,append,reshape等改变数据内存形状的操作后，需要使用numpy.ascontiguousarray对内存进行重新排序成连续的
        如使用非图像数据，也是转为numpy.ndarray数据类型再进行Tensor转换，使用{tensor_data} = Tensor({numpy_data})方式
        外部文件读入的numpy输入（例如np.fromfile）需要reshpe为对应的shape
        for i in range(input.shape[0]):
            input_tensor = Tensor(inputs[i, :].reshape(1,-1))
            input_tensor.to_device(device_id)
            input_tensors.append(input_tensor)

        """
    outputs = yolov3.infer(imageTensorList)
    yolov3_post = post.Yolov3PostProcess(
        config_path=config_path, label_path=label_path)  # 构造对应的后处理对象

    resizeInfo = base.ResizedImageInfo()
    resizeInfo.heightResize = yolo_resizelen
    resizeInfo.widthResize = yolo_resizelen
    resizeInfo.heightOriginal = decodedImg.original_height
    resizeInfo.widthOriginal = decodedImg.original_width

    inputs = []
    for i in range(len(outputs)):
        outputs[i].to_host()
        n = np.array(outputs[i])
        tensor = BTensor(n)  # 后处理需要使用baseTensor类型来构建，文档不全
        inputs.append(base.batch([tensor] * 2, keep_dims=True))

    results = yolov3_post.process(inputs, [resizeInfo] * 2)

    for i in range(len(results)):
        for j in range(len(results[i])):
            print(results[i][j].x0)
            print(results[i][j].y0)
            print(results[i][j].x1)
            print(results[i][j].y1)
            print(results[i][j].confidence)
            print(results[i][j].classId)
            print(results[i][j].className)


try:
    main()
except Exception as e:
    print(e)
